from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载预训练的模型和分词器
model_name = "microsoft/DialoGPT-medium"  # 或者选择其他你感兴趣的模型
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 定义了一个generate_response函数，它接受用户输入，通过模型生成回复，并返回该回复。
def generate_response(input_text):
    # 对输入文本进行编码
    inputs = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')
    
    # 生成回复
    outputs = model.generate(inputs, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    
    # 解码并输出生成的回复
    response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
    return response

if __name__ == "__main__":
    print("开始与AI代理对话（输入'exit'退出）：")
    while True:
        input_text = input("你: ")
        if input_text.lower() == 'exit':
            break
        response = generate_response(input_text)
        print(f"AI: {response}")